import torch
import sys

sys.path.append("/Users/zzq/Developer/python/SAM/reproduction/infini-transformer/")
# 将指定路径添加到系统路径中，以便导入自定义模块

from infini_transformer.compressive_memory import CompressiveMemory
from infini_transformer.positional_embeddings import RoPEEmbeddings

# 从infini_transformer模块中导入CompressiveMemory类
# 从infini_transformer.positional_embeddings模块中导入RoPEEmbeddings类

# 实例化 RoPEEmbeddings
rope_embedder = RoPEEmbeddings(dim=64, seq_len=65536)  # 确保参数与模型匹配

# 创建CompressiveMemory的实例，设置其参数
cm = CompressiveMemory(
    dim_input=768,  # 输入的维度
    dim_key=64,  # 键的维度
    dim_value=64,  # 值的维度
    num_heads=8,  # 注意力头的数量
    segment_len=2048,  # 段的长度
    sampling_factor=None,  # 采样因子
    update="linear",  # 更新策略
    causal=True,  # 是否为因果模型
    position_embedder=rope_embedder,  # 使用 RoPEEmbeddings 实例
    init_state_learnable=False,  # 初始状态是否可学习
)

batch = torch.randn(2, 65536, 768)  # 创建一个随机张量，形状为(2, 65536, 768)
# 其中2是批次大小，65536是序列长度，768是输入维度

output = cm(batch)  # 将批次数据传入CompressiveMemory实例，得到输出
print(output)
